from rest_framework import exceptions, status
from rest_framework.views import exception_handler

def custom_exception_handler(exc, context):
    response = exception_handler(exc, context)

    if response is not None:
        exc_class = exc.__class__

        if exc_class in [exceptions.ValidationError]:
            response.data = {
                'success': False,
                'message': response.data,
            }

            response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY

        if exc_class in [exceptions.AuthenticationFailed]:
            response.data = {
                'message': "Login failed",
            }
            response.status_code = status.HTTP_403_FORBIDDEN

        if exc_class in [exceptions.PermissionDenied]:
            response.data = {
                'message': "Forbidden for you",
            }
            response.status_code = status.HTTP_403_FORBIDDEN

        if exc_class in [exceptions.NotFound]:
            response.data = {
                'message': "Not found",
            }
            response.status_code = status.HTTP_404_NOT_FOUND

    return response

